package algorithm.graph.adjacencymatrix;

import algorithm.graph.base.Edge;
import algorithm.graph.base.GraphUtil;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

public class TopologicalSort {
    AdjacencyMatrixGraph graph;

    TopologicalSort(AdjacencyMatrixGraph graph) {
        this.graph = graph;
    }

    public static List<Integer> sort(AdjacencyMatrixGraph graph) {
        TopologicalSort topologicalSort = new TopologicalSort(graph);
        return topologicalSort.sort();
    }

    public List<Integer> sort() {
        List<Edge> edgeList = initEdgeList(graph);
        List<Integer> unvisitedNodeList = initNodeList(graph);
        List<Integer> visitedNodeList = new ArrayList<>();
        while (!unvisitedNodeList.isEmpty()) {
            Integer nodeId = findNextUnvisitedNode(unvisitedNodeList, edgeList);
            if (nodeId == null) {
                throw new RuntimeException("non directed acyclic graph");
            } else {
                GraphUtil.removeItemFromList(unvisitedNodeList, it -> it == nodeId);
                visitedNodeList.add(nodeId);
                GraphUtil.removeItemFromList(edgeList, it -> it.from() == nodeId);
            }
        }
        return visitedNodeList;
    }

    private Integer findNextUnvisitedNode(List<Integer> unvisitedNodeList, List<Edge> edgeList) {
        Map<Integer, List<Edge>> nodeEdgeListMap = edgeList.stream().collect(Collectors.groupingBy(Edge::to));
        for (Integer nodeId : unvisitedNodeList) {
            List<Edge> nodeEdgeList = nodeEdgeListMap.getOrDefault(nodeId, new ArrayList<>());
            if (nodeEdgeList.size() == 0) {
                return nodeId;
            }
        }
        return null;
    }

    private List<Integer> initNodeList(AdjacencyMatrixGraph graph) {
        return IntStream.range(0, graph.getNodeCount()).boxed().collect(Collectors.toCollection(ArrayList::new));
    }

    private List<Edge> initEdgeList(AdjacencyMatrixGraph graph) {
        return GraphUtil.getDirectedGraphEdgeList(graph);
    }
}
